import random
import numpy as np

from beam.util_psll import get_psll_by_phase


class RISGreyWolfOptimizer:
    __bit_num = 0  # 比特数
    __beam_num = 0  # 波束数

    def __init__(self, bit_num, beam_num, num_wolves=50, max_iter=100, a_decay=2):
        self.__bit_num = bit_num
        self.__beam_num = beam_num
        self.num_wolves = num_wolves
        self.max_iter = max_iter
        self.a_decay = a_decay  # 攻击参数衰减系数
        self.wolves = None
        self.alpha, self.beta, self.delta = None, None, None
        self.best_fitness = None
        self.best_fitness_history = []
        self.best_individual_history = []

    def fitness(self, phase):
        return get_psll_by_phase(phase, self.__bit_num, self.__beam_num)

    def apply_boundary_conditions(self, wolf):
        # 应用边界条件，确保新位置在合法范围内
        return np.clip(wolf, -180, 180)

    def initialize_pack(self, phase_mix_init):
        # 初始化第一个狼的位置为phase_mix_init
        self.wolves = [np.copy(phase_mix_init)]
        # 对于剩下的狼，基于phase_mix_init进行随机变异
        for _ in range(1, self.num_wolves):
            # 生成一个与phase_mix_init形状相同的布尔掩码，每个位置有10%的概率为True
            mutation_mask = np.random.rand(*phase_mix_init.shape) < 0.20
            # 生成随机偏移量
            random_offsets = np.random.uniform(-180, 180, phase_mix_init.shape)
            # 创建新的狼位置，并应用布尔掩码进行变异
            new_wolf = np.copy(phase_mix_init)
            new_wolf[mutation_mask] += random_offsets[mutation_mask]
            # 应用边界条件
            new_wolf = self.apply_boundary_conditions(new_wolf)
            # 将新位置添加到狼群中
            self.wolves.append(new_wolf)
        # 计算初始适应度并找到α、β、δ
        fitness_scores = [self.fitness(wolf) for wolf in self.wolves]
        sorted_indices = np.argsort(fitness_scores)
        self.alpha = self.wolves[sorted_indices[0]].copy()
        self.beta = self.wolves[sorted_indices[1]].copy()
        self.delta = self.wolves[sorted_indices[2]].copy()
        self.best_fitness = fitness_scores[sorted_indices[0]]
        self.best_fitness_history.append(self.best_fitness)
        self.best_individual_history.append(self.alpha)

    def update_position(self, iteration):
        a = 2 - iteration * ((2) / self.max_iter)  # 更新攻击参数a
        for i, wolf in enumerate(self.wolves):
            r1_alpha = np.random.rand(*wolf.shape)
            r2_alpha = np.random.rand(*wolf.shape)
            A_alpha = 2 * a * r1_alpha - a
            C_alpha = 2 * r2_alpha
            D_alpha = abs(C_alpha * self.alpha - wolf)
            X1 = self.alpha - A_alpha * D_alpha
            #
            r1_beta = np.random.rand(*wolf.shape)
            r2_beta = np.random.rand(*wolf.shape)
            A_beta = 2 * a * r1_beta - a
            C_beta = 2 * r2_beta
            D_beta = abs(C_beta * self.beta - wolf)
            X2 = self.beta - A_beta * D_beta
            #
            r1_delta = np.random.rand(*wolf.shape)
            r2_delta = np.random.rand(*wolf.shape)
            A_delta = 2 * a * r1_delta - a
            C_delta = 2 * r2_delta
            D_delta = abs(C_delta * self.delta - wolf)
            X3 = self.delta - A_delta * D_delta
            #
            new_wolf = (X1 + X2 + X3) / 3  # 更新位置
            # 应用边界条件
            new_wolf = self.apply_boundary_conditions(new_wolf)
            self.wolves[i] = new_wolf

    def run(self, phase_mix_init, logger):
        logger.info("num_wolves=%d, max_iter=%d, a_decay=%d" % (self.num_wolves, self.max_iter, self.a_decay))
        # 初始化返回值
        self.initialize_pack(phase_mix_init)
        self.best_individual = phase_mix_init
        self.best_fitness = self.fitness(phase_mix_init)
        self.best_fitness_history = []
        self.best_individual_history = []
        # 更新迭代
        for iteration in range(self.max_iter):
            # 更新狼群位置
            self.update_position(iteration)
            # 计算新的适应度值
            fitness_scores = [self.fitness(wolf) for wolf in self.wolves]
            sorted_indices = np.argsort(fitness_scores)
            # 更新α、β、δ
            self.alpha = self.wolves[sorted_indices[0]].copy()
            self.beta = self.wolves[sorted_indices[1]].copy()
            self.delta = self.wolves[sorted_indices[2]].copy()
            # 更新最佳适应度
            if fitness_scores[sorted_indices[0]] < self.best_fitness:
                self.best_fitness = fitness_scores[sorted_indices[0]]
                self.best_individual = self.alpha
            # 记录最佳适应度曲线
            self.best_fitness_history.append(self.best_fitness)
            self.best_individual_history.append(self.best_individual)
            logger.info("iteration=%d: self.best_fitness=%f" % (iteration, self.best_fitness))

        return self.best_individual, self.best_fitness, self.best_fitness_history, self.best_individual_history
